package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"regexp"
	"time"

	_ "github.com/go-sql-driver/mysql"
)

type ConceptItem struct {
	Code  string `json:"f12"` // Concept code
	Type  int    `json:"f13"` // Concept type (90)
	Name  string `json:"f14"` // Concept name
	Value int64  `json:"f62"` // Some numeric value
}

// Structs for concept block API response
type ConceptResponse struct {
	Rc   int `json:"rc"`
	Data struct {
		Total        int           `json:"total"`
		ConceptItems []ConceptItem `json:"diff"`
	} `json:"data"`
}

// Stock item struct
type StockItem struct {
	Code     string  `json:"f12"` // Stock code
	Name     string  `json:"f14"` // Stock name
	Market   int     `json:"f13"` // Market (sh:1, sz:0)
	Price    float64 //`json:"f2"`  // Current price,存在-
	Change   float64 //`json:"f3"` // Change percentage
	Volume   float64 //`json:"f5"` // Volume (手)
	Turnover float64 //`json:"f6"` // Turnover (万元)
}

// Struct for stock detail API response
type StockResponse struct {
	Rc   int `json:"rc"`
	Data struct {
		Total int         `json:"total"`
		Diff  []StockItem `json:"diff"`
	} `json:"data"`
}

// Struct for mapping concepts to stocks
type ConceptStock struct {
	ConceptCode string  `db:"bk_code"`
	StockCode   string  `db:"stock_code"`
	Price       float64 `db:"price"`
	Change      float64 `db:"change_percentage"`
	Volume      float64 `db:"volume"`
	Turnover    float64 `db:"turnover"`
}

func GetByUrl(url string) ([]byte, error) {
	referer := "https://data.eastmoney.com/bkzj/gn.html"

	client := &http.Client{}
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		fmt.Println("创建请求失败:", url, nil)
		return nil, err
	}
	req.Header.Add("Referer", referer)
	req.Header.Add("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36")
	resp, err := client.Do(req)
	if resp != nil {
		defer resp.Body.Close()
	}
	if err != nil {
		return nil, err
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	// 处理JSONP响应：提取JSON部分
	re := regexp.MustCompile(`\w+$([\s\S]*)$;?$`)
	matches := re.FindStringSubmatch(string(body))

	var jsonData string
	if len(matches) > 1 {
		jsonData = matches[1]
	} else {
		jsonData = string(body)
	}

	return []byte(jsonData), nil
}

// 初始化MySQL数据库连接
func initDB() (*sql.DB, error) {
	dsn := "root:root@tcp(127.0.0.1:3306)/stock?charset=utf8mb4&parseTime=True&loc=Local"
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return nil, err
	}

	// 测试连接
	err = db.Ping()
	return db, err
}

// 创建数据表
func createTables(db *sql.DB) error {
	// 创建概念板块表
	gnbkSQL := `
	CREATE TABLE IF NOT EXISTS gnbkv1 (
		code VARCHAR(20) PRIMARY KEY,
		name VARCHAR(100),
		value BIGINT,
		updated_at datetime DEFAULT CURRENT_TIMESTAMP
	)`
	_, err := db.Exec(gnbkSQL)
	if err != nil {
		return err
	}

	// 创建股票表
	stockSQL := `
	CREATE TABLE IF NOT EXISTS stockv1 (
		code VARCHAR(20) PRIMARY KEY,
		name VARCHAR(100),
		market INT,
		updated_at datetime DEFAULT CURRENT_TIMESTAMP
	)`
	_, err = db.Exec(stockSQL)
	if err != nil {
		return err
	}

	// 创建概念与股票映射表
	conceptStockSQL := `
	CREATE TABLE IF NOT EXISTS bk_stockv1 (
		bk_code VARCHAR(20),
		stock_code VARCHAR(20),
		price DECIMAL(10,2),
		change_percentage DECIMAL(5,2),
		volume DECIMAL(15,2),
		turnover DECIMAL(15,2),
		updated_at datetime DEFAULT CURRENT_TIMESTAMP,
		PRIMARY KEY (bk_code, stock_code)
	)`
	_, err = db.Exec(conceptStockSQL)
	return err
}

// 同步概念数据到数据库
func syncConcepts(db *sql.DB, concepts []ConceptItem) error {

	// 获取现有概念ID
	var existingCodes []string
	rows, err := db.Query("SELECT code FROM gnbkv1")
	if err != nil {
		return err
	}
	defer rows.Close()

	for _, concept := range concepts {
		existingCodes = append(existingCodes, concept.Code)
		// 插入或更新概念数据
		_, err = db.Exec(
			"INSERT INTO gnbkv1 (code, name, value) VALUES (?, ?, ?) ON DUPLICATE KEY UPDATE name = VALUES(name), value = VALUES(value), updated_at = CURRENT_TIMESTAMP",
			concept.Code,
			concept.Name,
			concept.Value,
		)
		if err != nil {
			return err
		}
	}

	// 删除不存在的概念
	if len(existingCodes) > 0 {
		// 构建查询参数
		placeholders := make([]string, len(existingCodes))
		args := make([]interface{}, len(existingCodes))
		for i, code := range existingCodes {
			placeholders[i] = "?"
			args[i] = code
		}

		// 查询不在新数据中的旧概念
		query := fmt.Sprintf("DELETE FROM gnbkv1 WHERE code NOT IN (%s)",
			stringJoin(placeholders, ","))

		_, err = db.Exec(query, args...)
		if err != nil {
			return err
		}
		// 删除不存在的概念股票关系
		query = fmt.Sprintf("DELETE FROM bk_stockv1 WHERE bk_code NOT IN (%s)",
			stringJoin(placeholders, ","))

		_, err = db.Exec(query, args...)
		if err != nil {
			return err
		}
	}

	// 提交事务
	return nil
}

// 同步股票数据到数据库
func syncStocks(db *sql.DB, stocks []StockItem) error {
	// 开始事务

	// 获取现有股票代码
	var existingCodes []string

	// 插入或更新股票数据
	for _, stock := range stocks {
		existingCodes = append(existingCodes, stock.Code)
		_, err := db.Exec(
			"INSERT INTO stockv1 (code, name, market) VALUES (?, ?, ?) ON DUPLICATE KEY UPDATE name = VALUES(name), market = VALUES(market), updated_at = CURRENT_TIMESTAMP",
			stock.Code,
			stock.Name,
			stock.Market,
		)
		if err != nil {
			return err
		}
	}

	// 删除不存在的股票
	if len(existingCodes) > 0 {
		// 构建查询参数
		placeholders := make([]string, len(existingCodes))
		args := make([]interface{}, len(existingCodes))
		for i, code := range existingCodes {
			placeholders[i] = "?"
			args[i] = code
		}

		// 查询不在新数据中的旧股票
		query := fmt.Sprintf("DELETE FROM stockv1 WHERE code NOT IN (%s)",
			stringJoin(placeholders, ","))

		_, err := db.Exec(query, args...)
		if err != nil {
			return err
		}
	}

	return nil
}

// 同步概念与股票的关系数据
func syncConceptStocks(db *sql.DB, conceptCode string, stocks []StockItem) error {

	// 删除现有的概念关联
	_, err := db.Exec("DELETE FROM bk_stockv1 WHERE bk_code = ?", conceptCode)
	if err != nil {
		return err
	}

	// 插入新的关联数据
	stmt, err := db.Prepare("INSERT INTO bk_stockv1 (bk_code, stock_code, price, change_percentage, volume, turnover) VALUES (?, ?, ?, ?, ?, ?)")
	if err != nil {
		return err
	}
	defer stmt.Close()

	for _, stock := range stocks {
		_, err = stmt.Exec(
			conceptCode,
			stock.Code,
			stock.Price,
			stock.Change,
			stock.Volume,
			stock.Turnover,
		)
		if err != nil {
			return err
		}
	}
	return nil
}

// 辅助函数：字符串连接
func stringJoin(elems []string, sep string) string {
	switch len(elems) {
	case 0:
		return ""
	case 1:
		return elems[0]
	default:
		n := len(sep) * (len(elems) - 1)
		for i := 0; i < len(elems); i++ {
			n += len(elems[i])
		}

		b := make([]byte, n)
		blen := copy(b, elems[0])
		for i := 1; i < len(elems); i++ {
			blen += copy(b[blen:], sep)
			blen += copy(b[blen:], elems[i])
		}
		return string(b)
	}
}

func main() {

	startTime := time.Now().UnixMilli()
	// 连接数据库
	db, err := initDB()
	if err != nil {
		fmt.Printf("Error connecting to database: %v\n", err)
		return
	}
	defer db.Close()

	// 创建数据表
	err = createTables(db)
	if err != nil {
		fmt.Printf("Error creating tables: %v\n", err)
		return
	}

	// 获取板块信息
	conceptAPI := "https://data.eastmoney.com/dataapi/bkzj/getbkzj?key=f62&code=m%3A90%2Bt%3A3"

	body, err := GetByUrl(conceptAPI)
	if err != nil {
		fmt.Printf("Error reading response: %v\n", err)
		return
	}

	// Parse concept block data
	var conceptResp ConceptResponse
	if err := json.Unmarshal(body, &conceptResp); err != nil {
		fmt.Printf("Error parsing concept blocks: %v\n", err)
		return
	}

	if conceptResp.Data.Total == 0 {
		fmt.Println("No concept blocks found")
		return
	}

	fmt.Printf("Found %d concept blocks\n", conceptResp.Data.Total)

	// 同步所有概念数据到数据库
	err = syncConcepts(db, conceptResp.Data.ConceptItems)
	if err != nil {
		fmt.Printf("Error syncing concepts: %v\n", err)
		return
	}

	// Process concepts sequentially to avoid rate limiting
	for _, concept := range conceptResp.Data.ConceptItems {
		fmt.Printf("Processing concept: %s\n", concept.Name)

		var allStocks []StockItem
		page := 1
		pageSize := 500
		totalStocks := 0

		for {
			// Fetch stock data for this concept with pagination
			stockAPI := fmt.Sprintf("https://push2.eastmoney.com/api/qt/clist/get?fid=f3&po=1&pz=%d&pn=%d&np=1&fltt=2&invt=2&fs=b:%s&fields=f12,f14,f2,f3,f5,f6,f13", pageSize, page, concept.Code)

			// Create request with headers
			req, err := http.NewRequest("GET", stockAPI, nil)
			if err != nil {
				fmt.Printf("Error creating request for %s page %d: %v\n", concept.Name, page, err)
				break
			}
			req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36")
			req.Header.Set("Accept", "application/json")
			req.Header.Set("Referer", "https://data.eastmoney.com/")
			req.Header.Set("Origin", "https://data.eastmoney.com")

			// Add random jitter
			// jitter := time.Duration(rand.Intn(1000)) * time.Millisecond
			// time.Sleep(jitter)

			// Make the request with retries
			maxRetries := 5
			var body []byte
			success := false
			for i := 0; i < maxRetries; i++ {
				resp, err := http.DefaultClient.Do(req)
				if err != nil {
					fmt.Printf("Error fetching stocks for %s page %d (attempt %d): %v\n", concept.Name, page, i+1, err)
					time.Sleep(time.Duration(i+1) * 3 * time.Second)
					continue
				}

				body, err = ioutil.ReadAll(resp.Body)
				resp.Body.Close()
				if err != nil {
					fmt.Printf("Error reading stock response for %s page %d (attempt %d): %v\n", concept.Name, page, i+1, err)
					time.Sleep(time.Duration(i+1) * 3 * time.Second)
					continue
				}

				// Check for HTML response (rate limiting)
				if len(body) > 0 && body[0] == '<' {
					fmt.Printf("HTML response received for %s page %d (attempt %d), likely rate limited\n", concept.Name, page, i+1)
					// Save HTML for debugging
					ioutil.WriteFile(fmt.Sprintf("error_%s_%d.html", concept.Code, i), body, 0644)
					time.Sleep(time.Duration(i+1) * 3 * time.Second)
					continue
				}

				success = true
				break
			}

			if !success {
				fmt.Printf("Skipping %s page %d after %d attempts due to failures\n", concept.Name, page, maxRetries)
				break
			}

			// Remove Byte Order Mark (BOM) if present
			if len(body) > 3 && body[0] == 0xEF && body[1] == 0xBB && body[2] == 0xBF {
				body = body[3:]
			}

			var stockResp StockResponse
			if err := json.Unmarshal(body, &stockResp); err != nil {
				fmt.Printf("Error parsing stocks for %s page %d: %v\n%s\n", concept.Name, page, err, string(body))
				break
			}

			// First page sets total stocks
			if page == 1 {
				totalStocks = stockResp.Data.Total
			}
			// if page > 11 {
			// 	break
			// }

			// Append stocks from this page
			allStocks = append(allStocks, stockResp.Data.Diff...)
			fmt.Printf("Page %d: Got %d stocks for %s (total: %d/%d)\n", page, len(stockResp.Data.Diff), concept.Name, len(allStocks), totalStocks)

			// Check if we have all stocks
			if len(allStocks) >= totalStocks || len(stockResp.Data.Diff) == 0 {
				break
			}

			page++

			// Sleep between pages to avoid rate limiting
			time.Sleep(1*time.Second + time.Duration(rand.Intn(100))*time.Millisecond)
		}

		// 同步股票数据到数据库
		err = syncStocks(db, allStocks)
		if err != nil {
			fmt.Printf("Error syncing stocks for %s: %v\n", concept.Name, err)
		}

		// 同步概念与股票的关系数据
		err = syncConceptStocks(db, concept.Code, allStocks)
		if err != nil {
			fmt.Printf("Error syncing concept stocks for %s: %v\n", concept.Name, err)
		}

		fmt.Printf("Completed %s: %d stocks\n", concept.Name, len(allStocks))

		// Sleep between concepts to avoid rate limiting
		time.Sleep(1*time.Second + time.Duration(rand.Intn(1000))*time.Millisecond)
	}

	useTime := time.Now().UnixMilli() - startTime
	fmt.Println("Successfully saved concept and stock data to MySQL database. use time:", useTime)
}
